{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a5cfd7e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib.patches import Patch\n",
    "\n",
    "from sklearn.model_selection import(\n",
    "    GroupKFold,\n",
    "    GroupShuffleSplit,\n",
    "    KFold,\n",
    "    ShuffleSplit,\n",
    "    StratifiedGroupKFold,\n",
    "    StratifiedKFold,\n",
    "    StratifiedShuffleSplit,\n",
    "    TimeSeriesSplit,\n",
    ")\n",
    "\n",
    "rng = np.random.RandomState(1338)\n",
    "cmap_data = plt.cm.Paired\n",
    "cmap_cv = plt.cm.coolwarm\n",
    "n_splits = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ad16e724",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0, 0.1), (1, 0.3), (2, 0.6)]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(enumerate(percentiles_classes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6982c0e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "xx = [[ii] * int(100 * perc) for ii, perc in enumerate(percentiles_classes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fa1453d0",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "y=np.hstack([[ii] * int(100 * perc) for ii, perc in enumerate(percentiles_classes)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "13a307f5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "       1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2,\n",
       "       2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,\n",
       "       2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,\n",
       "       2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "17b55023",
   "metadata": {},
   "outputs": [],
   "source": [
    "group_prior = rng.dirichlet([2] * 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f66647c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.09120656, 0.0746941 , 0.07980615, 0.03136101, 0.17038876,\n",
       "       0.15518784, 0.04081938, 0.09239807, 0.1820129 , 0.08212523])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "group_prior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "6a3bf6a4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 7, 12,  9,  3, 16, 11,  5,  8, 21,  8])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rng.multinomial(100, group_prior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "2cdaa656",
   "metadata": {},
   "outputs": [],
   "source": [
    "groups = np.repeat(np.arange(10),rng.multinomial(100, group_prior))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "ca16c048",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3,\n",
       "       3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5,\n",
       "       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,\n",
       "       5, 5, 5, 5, 6, 6, 7, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,\n",
       "       8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08f52e63",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "695ff1d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the class/group data\n",
    "n_points = 100\n",
    "X = rng.randn(100, 10)\n",
    "\n",
    "percentiles_classes = [0.1, 0.3, 0.6]\n",
    "y=np.hstack([[ii] * int(100 * perc) for ii, perc in enumerate(percentiles_classes)])\n",
    "\n",
    "# Generate uneven groups\n",
    "group_prior = rng.dirichlet([2] * 10)\n",
    "groups = np.repeat(np.arange(10),rng.multinomial(100, group_prior))\n",
    "\n",
    "def visualize_groups(classes, groups, name):\n",
    "    # visualize dataset groups\n",
    "    fig, ax = plt.subplots()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9b5c42b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
